from typing import List, Dict
from schedule.PowerSchedule import PowerSchedule
from utils.Seed import Seed


class CrashPowerSchedule(PowerSchedule):
    def __init__(self, exponent: float = 1.0) -> None:
        super().__init__()
        self.crash_counts: Dict[str, int] = {}
        self.exponent = exponent

    def assign_energy(self, population: List[Seed]) -> None:
        """Assign higher energy to seeds that have triggered more crashes"""
        for seed in population:
            seed_data = seed.data
            crash_count = self.crash_counts.get(seed_data, 0)
            # 加 1 是为了防止 crash_count 为 0 导致种子能量为 0
            seed.energy = (crash_count + 1) ** self.exponent

    def increment_crash_count(self, seed: Seed) -> None:
        """Increment the crash count for the given seed"""
        seed_data = seed.data
        if seed_data is not None:
            self.crash_counts[seed_data] = self.crash_counts.get(seed_data, 0) + 1
